import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
import os
from scipy import interpolate


def draw2DFigure(probe_data_obj, figure_conf):
    img_dict = {}
    plt.rcParams['font.sans-serif']=['SimHei']      # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False      # 用来正常显示负号
    for key in probe_data_obj:
        df = probe_data_obj[key]
        df.replace('None',np.nan,inplace=True)
        df['GumResult'] = df['GumResult'].astype('float64')
        df = df.pivot_table(index='Height',columns='DataStartTime',values='GumResult')
        fig, ax = plt.subplots()
        X, Y = np.meshgrid(np.arange(len(df.columns)), df.index)
        c = ax.contour(X, Y, df, levels=8, colors="black", linewidths=0.5)
        ax.clabel(c, inline=True, fontsize=10)
        c = ax.contourf(X, Y, df, levels=40, cmap='jet')
        plt.colorbar(c)
        ax.set_xlabel('时间')
        ax.set_ylabel('高度(km)')
        if key == "T":
            ax.set_title('温度测量不确定度时空分布(K)')
        elif key == "D":
            ax.set_title('大气密度测量不确定度时空分布(%)')
        elif key == "N":
            ax.set_title('钠密度测量不确定度时空分布(%)')
        elif key == "F":
            ax.set_title('铁密度测量不确定度时空分布(%)')
        elif key == "W1":
            ax.set_title('纬向风测量不确定度时空分布(m/s)')
        elif key == "W2":
            ax.set_title('经向风测量不确定度时空分布(m/s)')
        elif key == "W":
            ax.set_title('风速测量不确定度时空分布(m/s)')
        elif key == "X":
            ax.set_title('风向测量不确定度时空分布(°)')
        fig_path = os.path.join(figure_conf["figure_location"], key+'.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key] = fig_path
    return img_dict

def dataCompareDrawIndexFigure(probe_data_obj, figure_conf):
    img_dict = {}
    plt.rcParams['font.sans-serif']=['SimHei']      # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False      # 用来正常显示负号
    for key in probe_data_obj:
        img_dict[key] = {}
        df1 = probe_data_obj[key][0]
        df1.replace('None',np.nan,inplace=True)
        df1['Value'] = df1['Value'].astype('float64')
        df2 = probe_data_obj[key][-1]
        df2.replace('None',np.nan,inplace=True)
        df2['Value'] = df2['Value'].astype('float64')
        df1 = df1.pivot_table(index='Height',columns='DataStartTime',values='Value')
        df2 = df2.pivot_table(index='Height',columns='DataStartTime',values='Value')
        common_time = set(df1.columns)&set(df2.columns)
        df1 = df1[common_time]
        df2 = df2[common_time]
        height = df1.index.tolist()+df2.index.tolist()
        height_min = np.floor(np.min(height))
        height_max = np.ceil(np.max(height))
        height = np.arange(height_min,height_max)
        f = interpolate.interp1d(df1.index,df1.values.T)
        arr1 = f(height).T
        f = interpolate.interp1d(df2.index,df2.values.T)
        arr2 = f(height).T
        if key in ['D','N','F']:
            arr_dev = (arr2-arr1)/arr1*100
        else:
            arr_dev = arr2-arr1
        arr_MD = np.nanmean(arr_dev,axis=1)
        arr_SD = np.nanstd(arr_dev,axis=1)
        arr1 = arr1.reshape(-1)
        arr2 = arr2.reshape(-1)
        COEF = np.corrcoef(arr1, arr2)[0,1]
        if key == "T":
            xlabel = '温度(K)'
        elif key == "D":
            xlabel = '大气密度(%)'
        elif key == "N":
            xlabel = '钠密度(%)'
        elif key == "F":
            xlabel = '铁密度(%)'
        elif key == "W1":
            xlabel = '纬向风(m/s)'
        elif key == "W2":
            xlabel = '经向风(m/s)'
        elif key == "W":
            xlabel = '风速(m/s)'
        elif key == "X":
            xlabel = '风向(°)'
        fig, ax = plt.subplots()
        ax.plot(arr_MD,height)
        ax.set_xlabel(xlabel)
        ax.set_ylabel('高度(km)')
        fig_path = os.path.join(figure_conf["figure_location"], 'MD.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key]['MD'] = fig_path
        fig, ax = plt.subplots()
        ax.plot(arr_SD,height)
        ax.set_xlabel(xlabel)
        ax.set_ylabel('高度(km)')
        fig_path = os.path.join(figure_conf["figure_location"], 'SD.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key]['SD'] = fig_path
        fig, ax = plt.subplots()
        ax.scatter(arr1,arr2,s=5)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(xlabel)
        ax.legend(['相关系数={:.4f}'.format(COEF)])
        par = np.polyfit(arr1, arr2, 1)
        f = np.poly1d(par)
        ax.plot(arr1, f(arr1), '--')
        fig_path = os.path.join(figure_conf["figure_location"], 'COEF.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key]['COEF'] = fig_path
    return img_dict

def dataAnalysisDrawSpaceTime(probe_data_obj, figure_conf):
    img_dict = {}
    plt.rcParams['font.sans-serif']=['SimHei']      # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False      # 用来正常显示负号
    for key in probe_data_obj:
        df = probe_data_obj[key]
        df.replace('None',np.nan,inplace=True)
        df['Value'] = df['Value'].astype('float64')
        df = df.pivot_table(index='Height',columns='DataStartTime',values='Value')
        fig, ax = plt.subplots()
        X, Y = np.meshgrid(np.arange(len(df.columns)), df.index)
        c = ax.contour(X, Y, df, levels=8, colors="black", linewidths=0.5)
        ax.clabel(c, inline=True, fontsize=10)
        c = ax.contourf(X, Y, df, levels=40, cmap='jet')
        plt.colorbar(c)
        ax.set_xlabel('时间')
        ax.set_ylabel('高度(km)')
        if key == "T":
            ax.set_title('温度时空分布(K)')
        elif key == "D":
            ax.set_title('大气密度时空分布(%)')
        elif key == "N":
            ax.set_title('钠密度时空分布(%)')
        elif key == "F":
            ax.set_title('铁密度时空分布(%)')
        elif key == "W1":
            ax.set_title('纬向风时空分布(m/s)')
        elif key == "W2":
            ax.set_title('经向风时空分布(m/s)')
        elif key == "W":
            ax.set_title('风速时空分布(m/s)')
        elif key == "X":
            ax.set_title('风向时空分布(°)')
        fig_path = os.path.join(figure_conf["figure_location"], key+'.png')
        fig.savefig(fig_path, dpi=300, bbox_inches="tight")
        img_dict[key] = fig_path
    return img_dict

if __name__ == '__main__':
    # plt.rcParams['font.sans-serif']=['SimHei']      # 用来正常显示中文标签
    # plt.rcParams['axes.unicode_minus'] = False      # 用来正常显示负号
    # arr1 = np.arange(100)
    # arr2 = arr1**2
    # arr1 = arr1.reshape(-1)
    # arr2 = arr2.reshape(-1)
    # COEF = np.corrcoef(arr1, arr2)[0,1]
    # print(COEF)
    # fig, ax = plt.subplots()
    # ax.scatter(arr1,arr2,s=5)
    # ax.set_ylabel('高度(km)')
    # ax.legend(['相关系数={:.4f}'.format(COEF)])
    # par = np.polyfit(arr1, arr2, 1)
    # f = np.poly1d(par)
    # ax.plot(arr1, f(arr1), '--')
    # filepath = r'C:\Users\Administrator\Desktop\output_T.csv'
    filepath = r'D:\GBfiles\mycode\Workspace\MeasurementProject\measurement_python\output_T.csv'
    df = pd.read_csv(filepath,index_col=0)
    probe_data_obj = {'T':df}
    figure_conf = {"figure_location":r'D:\GBfiles\mycode\Workspace\MeasurementProject\measurement_python'}
    print("画二维彩图参数：", probe_data_obj)
    draw2DFigure(probe_data_obj, figure_conf)

